# encoding: utf-8
# @File  : tasks.py
# @Author: shaoyun
# @Date  :  2025/05/11
from typing import Optional, Dict, Any

from fastapi import APIRouter, Depends, Query, Path
from sqlalchemy.orm import Session
from starlette.status import HTTP_201_CREATED

from app.api.dependencies import get_current_user, get_db
from app.core.response import get_response
from app.core.status_codes import SUCCESS, ERROR,TASK_CREATION_FAILED, TASK_NOT_FOUND, PERMISSION_DENIED, INVALID_TASK_STATUS
from app.models.task import TaskStatus, Task
from app.models.user import User
from app.schemas.task import TaskCreate, TaskUpdate
from app.services.ai_optimizer import ai_optimizer
from app.services.spider_runner import SpiderRunner
from app.services.task_scheduler import TaskScheduler

router = APIRouter()


@router.post("/create", response_model=Dict[str, Any], status_code=HTTP_201_CREATED)
async def create_task(
        task_in: TaskCreate,
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user)
):
    """创建新任务"""
    try:
        # 预测任务优先级
        priority = await ai_optimizer.predict_task_priority(
            task_id="temp",
            spider_id=task_in.spider_id,
            parameters=task_in.parameters
        )

        # 如果AI预测的优先级比用户设置的高，则使用AI预测的优先级
        if priority > task_in.priority:
            task_in.priority = priority

        # 创建任务
        task = await TaskScheduler.create_task(
            db=db,
            user_id=current_user.id,
            spider_id=task_in.spider_id,
            name=task_in.name,
            parameters=task_in.parameters,
            priority=task_in.priority,
            cron=task_in.cron
        )

        # 将DB模型转换为字典
        task_data = {
            "id": task.id,
            "user_id": task.user_id,
            "spider_id": task.spider_id,
            "name": task.name,
            "parameters": task.parameters,
            "status": task.status.value,
            "priority": task.priority,
            "cron": task.cron,
            "created_at": task.created_at,
            "updated_at": task.updated_at,
            "started_at": task.started_at,
            "finished_at": task.finished_at
        }

        return get_response(SUCCESS, data=task_data)
    except Exception:
        return get_response(TASK_CREATION_FAILED, data=None)


@router.get("/list", response_model=Dict[str, Any])
async def get_tasks(
        skip: int = Query(0, ge=0),
        limit: int = Query(100, ge=1, le=1000),
        status: Optional[str] = None,
        spider_id: Optional[str] = None,
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user)
):
    """获取任务列表"""
    try:
        # 如果是管理员，可以查看所有任务，否则只能查看自己的任务
        user_id = None if current_user.role.value == "admin" else current_user.id

        tasks = await TaskScheduler.get_tasks(
            db=db,
            user_id=user_id,
            status=status,
            skip=skip,
            limit=limit
        )

        # 将DB模型转换为字典列表
        task_list = [
            {
                "id": task.id,
                "user_id": task.user_id,
                "spider_id": task.spider_id,
                "name": task.name,
                "parameters": task.parameters,
                "status": task.status.value,
                "priority": task.priority,
                "cron": task.cron,
                "created_at": task.created_at,
                "updated_at": task.updated_at,
                "started_at": task.started_at,
                "finished_at": task.finished_at
            }
            for task in tasks
        ]

        # 获取总数
        total_query = db.query(Task)
        if user_id:
            total_query = total_query.filter(Task.user_id == user_id)
        if status:
            try:
                task_status = TaskStatus(status)
                total_query = total_query.filter(Task.status == task_status)
            except ValueError:
                pass
        if spider_id:
            total_query = total_query.filter(Task.spider_id == spider_id)

        total = total_query.count()

        return get_response(SUCCESS, data={
            "items": task_list,
            "total": total,
            "page": (skip // limit) + 1,
            "limit": limit
        })
    except Exception:
        return get_response(ERROR,data=None)


@router.get("/{task_id}/detail", response_model=Dict[str, Any])
async def get_task(
        task_id: str = Path(...),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user)
):
    """获取任务详情"""
    try:
        # 使用工具函数获取任务并检查权限
        task = await get_task_and_check_permissions(task_id, db, current_user)

        # 获取任务结果
        results = await SpiderRunner.get_spider_results(task_id)

        # 构建响应数据
        task_data = {
            "id": task.id,
            "user_id": task.user_id,
            "spider_id": task.spider_id,
            "name": task.name,
            "parameters": task.parameters,
            "status": task.status.value,
            "priority": task.priority,
            "cron": task.cron,
            "created_at": task.created_at,
            "updated_at": task.updated_at,
            "started_at": task.started_at,
            "finished_at": task.finished_at,
            "results": results.get("data", [])
        }

        return get_response(SUCCESS, data=task_data)
    except ValueError as e:
        return get_response(e.args[0], data=None)
    except Exception:
        return get_response(ERROR, data=None)


@router.put("/{task_id}/update", response_model=Dict[str, Any])
async def update_task(
        task_update: TaskUpdate,
        task_id: str = Path(...),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user)
):
    """更新任务"""
    try:
        task = await TaskScheduler.get_task(db=db, task_id=task_id)

        if not task:
            return get_response(TASK_NOT_FOUND, data=None)

        # 检查权限
        if current_user.role.value != "admin" and task.user_id != current_user.id:
            return get_response(PERMISSION_DENIED, data=None)

        # 更新任务字段
        if task_update.name is not None:
            task.name = task_update.name

        if task_update.parameters is not None:
            task.parameters = task_update.parameters

        if task_update.priority is not None:
            task.priority = task_update.priority

        if task_update.cron is not None:
            task.cron = task_update.cron

        if task_update.status is not None:
            try:
                new_status = TaskStatus(task_update.status)
                await TaskScheduler.update_task_status(db, task_id, new_status)
            except ValueError:
                return get_response(INVALID_TASK_STATUS, data=None)

        db.commit()
        db.refresh(task)

        # 构建响应数据
        task_data = {
            "id": task.id,
            "user_id": task.user_id,
            "spider_id": task.spider_id,
            "name": task.name,
            "parameters": task.parameters,
            "status": task.status.value,
            "priority": task.priority,
            "cron": task.cron,
            "created_at": task.created_at,
            "updated_at": task.updated_at,
            "started_at": task.started_at,
            "finished_at": task.finished_at
        }

        return get_response(SUCCESS, data=task_data)
    except Exception:
        return get_response(ERROR, data=None)


@router.put("/{task_id}/restart", response_model=Dict[str, Any])
async def restart_task(
        task_id: str = Path(...),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user)
):
    """重启任务"""
    try:
        task = await TaskScheduler.get_task(db=db, task_id=task_id)

        if not task:
            return get_response(TASK_NOT_FOUND, data=None)

        # 检查权限
        if current_user.role.value != "admin" and task.user_id != current_user.id:
            return get_response(PERMISSION_DENIED, data=None)

        # 只有失败或取消的任务可以重启
        if task.status not in [TaskStatus.FAILED, TaskStatus.CANCELED]:
            return get_response(INVALID_TASK_STATUS, data={"message": "只有失败或已取消的任务可以重启"})

        # 重启任务：重置任务状态和时间字段
        task.status = TaskStatus.PENDING
        task.started_at = None
        task.finished_at = None

        # 如果有重试计数，可以重置
        # task.retry_count = 0

        db.commit()
        db.refresh(task)

        # 构建响应数据
        task_data = {
            "id": task.id,
            "user_id": task.user_id,
            "spider_id": task.spider_id,
            "name": task.name,
            "parameters": task.parameters,
            "status": task.status.value,
            "priority": task.priority,
            "cron": task.cron,
            "created_at": task.created_at,
            "updated_at": task.updated_at,
            "started_at": task.started_at,
            "finished_at": task.finished_at
        }

        return get_response(SUCCESS, data=task_data)
    except Exception:
        return get_response(ERROR, data=None)


@router.delete("/{task_id}/cancel", response_model=Dict[str, Any])
async def cancel_task(
        task_id: str = Path(...),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user)
):
    """取消任务"""
    try:
        task = await TaskScheduler.get_task(db=db, task_id=task_id)

        if not task:
            return get_response(TASK_NOT_FOUND, data=None)

        # 检查权限
        if current_user.role.value != "admin" and task.user_id != current_user.id:
            return get_response(PERMISSION_DENIED, data=None)

        # 只有待执行的任务可以取消
        if task.status != TaskStatus.PENDING:
            return get_response(INVALID_TASK_STATUS, data=None)

        success = await TaskScheduler.cancel_task(db, task_id)

        if not success:
            return get_response(ERROR, data=None)

        return get_response(SUCCESS, data=None)
    except Exception:
        return get_response(ERROR, data=None)


@router.get("/{task_id}/results", response_model=Dict[str, Any])
async def get_task_results(
        task_id: str = Path(...),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user)
):
    """获取任务结果"""
    try:
        task = await TaskScheduler.get_task(db=db, task_id=task_id)

        if not task:
            return get_response(TASK_NOT_FOUND, data=None)

        # 检查权限
        if current_user.role.value != "admin" and task.user_id != current_user.id:
            return get_response(PERMISSION_DENIED, data=None)

        # 获取任务结果
        results = await SpiderRunner.get_spider_results(task_id)

        return get_response(SUCCESS, data={
            "task_id": task_id,
            "status": task.status.value,
            "data": results.get("data", []),
            "created_at": results.get("created_at")
        })
    except Exception:
        return get_response(ERROR, data=None)


@router.get("/{task_id}/log", response_model=Dict[str, Any])
async def get_task_log(
        task_id: str = Path(...),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user)
):
    """获取任务执行日志"""
    try:
        task = await TaskScheduler.get_task(db=db, task_id=task_id)

        if not task:
            return get_response(TASK_NOT_FOUND, data=None)

        # 检查权限
        if current_user.role.value != "admin" and task.user_id != current_user.id:
            return get_response(PERMISSION_DENIED, data=None)

        # 获取任务执行日志
        log = await SpiderRunner.get_execution_log(task_id)

        return get_response(SUCCESS, data={
            "task_id": task_id,
            "success": log.get("success", False),
            "message": log.get("message", ""),
            "log": log.get("log", ""),
            "timestamp": log.get("timestamp")
        })
    except Exception:
        return get_response(ERROR, data=None)


async def get_task_and_check_permissions(
        task_id: str,
        db: Session,
        current_user: User
) -> Task:
    """
    获取任务并检查权限
    :param task_id: 任务ID
    :param db: 数据库会话
    :param current_user: 当前用户
    :return: Task 对象
    """
    # 获取任务
    task = await TaskScheduler.get_task(db=db, task_id=task_id)

    if not task:
        raise ValueError(TASK_NOT_FOUND)

    # 检查权限
    if current_user.role.value != "admin" and task.user_id != current_user.id:
        raise ValueError(PERMISSION_DENIED)

    return task